import time
import argparse
import psutil
import subprocess

pos = 0


def restart(file, idx, args):
    print("Restart!")
    if idx > 0 and file != "":
        # f = open("ignore_function_conv_as.txt", "a+")
        f = open(args.ignore_config, "a+")
        f.write("{}:{}".format(idx, file))
        f.close()

    proc = subprocess.Popen('python ' + 'generate_trace.py ' + '-e 1 -p {} '.format(args.data) +
                            '--trace_data {} '.format(args.trace_data) + '-i {} -nw 1 '.format(args.ignore_config) +
                            '> {}'.format(args.log),
                            shell=True)
    # proc = subprocess.Popen('python ' + 'generate_trace_as.py ' + '-e 1 -p tvm_data/data_conv/elf/ ' +
    #                         '--trace_data as_trace_data_all.pkl ' + '-i ignore_function_conv_as.txt ' +
    #                         '> gen_trace_as.log',
    #                         shell=True)
    return proc.pid


def check(pid, args):
    ignore_file = ""
    ignore_idx = 0
    try:
        print("Start checking...")
        fd = open(args.log, "r")

        for s in fd.readlines():
            # print(s)
            if "$" in s:
                ignore_file = s.split("$")[-1].split("?")[1]
                ignore_idx = int(s.split("$")[-1].split("?")[0])
        fd.close()
    except Exception:
        pass
    finally:
        fd.close()

    try:
        p = psutil.Process(pid)
        status = p.status()
        if status == "zombie":
            pid = restart(ignore_file, ignore_idx, args)
        return pid
    except psutil.NoSuchProcess:
        time.sleep(5)
        pid = restart(ignore_file, ignore_idx, args)
        return pid


def run(args):
    pid = args.pid
    while True:
        time.sleep(20)
        pid = check(pid, args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Generate the trace from CFG")
    parser.add_argument("-p", "--pid", type=int, help="The pid of process")
    parser.add_argument("-l", "--log", type=str, help="The name of the log file")
    parser.add_argument("-d", "--data", type=str, help="The path of .so data")
    parser.add_argument("-e", "--extend", type=int, default=0, help="The trigger of extending the graph")
    parser.add_argument("-t", "--trace_number", type=int, default=10, help="The number of traces")
    parser.add_argument("-i",
                        "--ignore_config",
                        type=str,
                        default="ignore_function.txt",
                        help="The config for ignoring files")
    parser.add_argument("-nw", "--no_w2v", default=0, type=int)
    parser.add_argument("--trace_data", type=str, default="None", help="The path of trace data")
    args = parser.parse_args()
    run(args)
